End-to-end local RAG on a Databricks T4 GPU

This quickstart composes a complete local retrieval-augmented generation (RAG) pipeline with Apache Spark and SynapseML:

1. encode document text and a question with a Hugging Face sentence transformer on a GPU worker;
2. rank the document embeddings with exact cosine similarity on that GPU; and
3. generate a grounded answer with HuggingFaceCausalLM and Phi-4-mini on the GPU.

Unlike the service-backed PDF Q&A quickstart, this example needs no model-service keys or vector database. For large corpora, replace the exact retrieval step with the indexed approach in the GPU approximate KNN quickstart.

The SynapseML GPU smoke test runs this notebook on Databricks Runtime 14.3 LTS ML with one Standard_NC16as_T4_v3 worker. PyTorch and CUDA come from that GPU runtime; do not install a separate CUDA toolkit or TensorRT-LLM. The Python libraries used by the notebook are pinned to the versions tested by the repository:

• transformers==4.49.0
• huggingface-hub==0.26.0
• sentence-transformers==4.0.2
• accelerate==0.26.0

The configuration below also pins each Hugging Face repository to an immutable commit SHA. The exact huggingface-hub version is pinned because it resolves those snapshots. Phi's model and tokenizer load from the same pinned local snapshot, and Transformers' native Phi implementation is used with remote model code disabled.

When these libraries are not already installed as cluster libraries, install the same versions and restart Python before continuing:

# %pip install transformers==4.49.0 huggingface-hub==0.26.0 sentence-transformers==4.0.2 accelerate==0.26.0
# dbutils.library.restartPython()

1. Configure the reproducible smoke path

Pull-request validation supplies the synapseml_ci_smoke widget. The smoke path still performs every GPU stage; it only uses fewer documents and fewer generated tokens. Model revisions are full Hugging Face commit SHAs so a future change to either repository cannot alter this example silently.

EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
EMBEDDING_REVISION = "1110a243fdf4706b3f48f1d95db1a4f5529b4d41"
GENERATION_MODEL = "microsoft/Phi-4-mini-instruct"
GENERATION_REVISION = "cfbefacb99257ffa30c83adab238a50856ac3083"
ANSWER_DOCUMENT_ID = "earth-view"

ci_smoke = False
if "dbutils" in globals():
    dbutils.widgets.text("synapseml_ci_smoke", "false")
    ci_smoke = dbutils.widgets.get("synapseml_ci_smoke").lower() == "true"

generation_tokens = 24 if ci_smoke else 64
retrieve_k = 2

2. Verify the Spark executor has a GPU

The Databricks test cluster intentionally uses a CPU driver and a T4 GPU worker. Probe the executor rather than the driver so the check matches where Spark inference runs.

import numpy as np
import pandas as pd
import pyspark.sql.functions as F
from pyspark.ml.functions import predict_batch_udf
from pyspark.sql.functions import pandas_udf
from pyspark.sql.types import ArrayType, FloatType, StringType


@pandas_udf(StringType())
def cuda_device_name(values: pd.Series) -> pd.Series:
    import torch

    if not torch.cuda.is_available():
        raise RuntimeError("This quickstart requires a CUDA-enabled Spark worker.")
    return pd.Series([torch.cuda.get_device_name(0)] * len(values))


gpu_name = (
    spark.range(1)
    .repartition(1)
    .select(cuda_device_name(F.col("id")).alias("gpu"))
    .first()["gpu"]
)
assert gpu_name, "The Spark worker did not report a CUDA device."
print(f"Spark executor GPU: {gpu_name}")

3. Create a small local knowledge base

The compact corpus keeps the tutorial deterministic and free of service credentials. In an application, replace this DataFrame with text extracted from PDFs using the preprocessing steps in the PDF Q&A quickstart.

documents = [
    (
        "earth-at-night",
        "Earth at night reveals cities and transportation networks through "
        "patterns of artificial light.",
    ),
    (
        "mars",
        "Mars is often called the red planet because iron minerals in its soil "
        "oxidize and appear red.",
    ),
    (
        "earth-view",
        "Apollo 14 astronaut Edgar Mitchell described Earth from space as "
        "a sparkling blue and white jewel.",
    ),
    (
        "earth-oceans",
        "Oceans cover most of Earth's surface and strongly influence weather "
        "and climate.",
    ),
]

if not ci_smoke:
    documents.extend(
        [
            (
                "earth-atmosphere",
                "Earth's atmosphere scatters blue light and protects life from "
                "much of the Sun's harmful radiation.",
            ),
            (
                "moon",
                "The Moon is Earth's only natural satellite and stabilizes the "
                "planet's axial wobble.",
            ),
        ]
    )

indexed_documents = [
    (input_position, document_id, text)
    for input_position, (document_id, text) in enumerate(documents)
]
answer_position = next(
    input_position
    for input_position, document_id, _ in indexed_documents
    if document_id == ANSWER_DOCUMENT_ID
)
assert answer_position >= retrieve_k

documents_df = spark.createDataFrame(
    indexed_documents, ["input_position", "document_id", "text"]
).repartition(1)
question = "What did astronaut Edgar Mitchell call Earth?"
assert documents_df.count() >= retrieve_k

4. Generate normalized sentence embeddings on the GPU

predict_batch_udf loads the model once per Python worker and batches Spark rows. The factory checks CUDA inside the worker, avoiding assumptions about the CPU driver. The revision applies to the complete Sentence Transformers snapshot, including its tokenizer, and remote model code stays disabled.

def make_sentence_embedder():
    import torch
    from sentence_transformers import SentenceTransformer

    if not torch.cuda.is_available():
        raise RuntimeError("Sentence embedding requires a CUDA-enabled Spark worker.")

    model = SentenceTransformer(
        EMBEDDING_MODEL,
        device="cuda",
        revision=EMBEDDING_REVISION,
        trust_remote_code=False,
    )

    def predict(text_batch):
        return model.encode(
            text_batch.tolist(),
            batch_size=32,
            convert_to_numpy=True,
            normalize_embeddings=True,
            show_progress_bar=False,
        )

    return predict


embed = predict_batch_udf(
    make_sentence_embedder,
    return_type=ArrayType(FloatType()),
    batch_size=32,
)

query_df = spark.createDataFrame(
    [(-1, "question", question)], ["input_position", "document_id", "text"]
)
texts_to_embed = documents_df.unionByName(query_df).repartition(1)
embedded_df = texts_to_embed.withColumn("embedding", embed(F.col("text"))).cache()

embedding_sizes = {
    row["embedding_size"]
    for row in embedded_df.select(F.size("embedding").alias("embedding_size"))
    .distinct()
    .collect()
}
assert embedding_sizes == {384}, f"Unexpected embedding sizes: {embedding_sizes}"

5. Retrieve context with exact cosine similarity on the GPU

For this tutorial-sized corpus, exact scoring is easier to understand and validate than an approximate index. Both embeddings are already normalized, but cosine similarity keeps the retrieval step explicit. A persisted input_position puts the answer document outside the first retrieve_k corpus rows. The smoke checks compare Spark's top-k with an independent Python sort of every scored row, require the answer to rank first with a strictly higher finite score, and prove that taking the first input rows would miss it. The GPU KNN quickstarts show scalable indexed alternatives.

@pandas_udf(FloatType())
def gpu_cosine_similarity(
    document_embeddings: pd.Series, query_embeddings: pd.Series
) -> pd.Series:
    import torch

    if not torch.cuda.is_available():
        raise RuntimeError("Similarity scoring requires a CUDA-enabled Spark worker.")

    document_tensor = torch.as_tensor(
        np.stack(document_embeddings.to_list()), dtype=torch.float32, device="cuda"
    )
    query_tensor = torch.as_tensor(
        np.stack(query_embeddings.to_list()), dtype=torch.float32, device="cuda"
    )
    scores = torch.nn.functional.cosine_similarity(document_tensor, query_tensor, dim=1)
    return pd.Series(scores.detach().cpu().numpy())


document_embeddings = embedded_df.filter(F.col("document_id") != "question")
query_embedding = embedded_df.filter(F.col("document_id") == "question").select(
    F.col("embedding").alias("query_embedding")
)

scored_df = (
    document_embeddings.crossJoin(query_embedding)
    .repartition(1)
    .withColumn(
        "similarity",
        gpu_cosine_similarity(F.col("embedding"), F.col("query_embedding")),
    )
    .cache()
)
ranked_df = scored_df.orderBy(F.desc("similarity"), F.asc("input_position"))

scored_rows = scored_df.select(
    "input_position", "document_id", "text", "similarity"
).collect()
all_similarities = [row["similarity"] for row in scored_rows]
assert all(np.isfinite(all_similarities)), all_similarities

expected_rows = sorted(
    scored_rows, key=lambda row: (-row["similarity"], row["input_position"])
)
retrieved_rows = (
    ranked_df.select("input_position", "document_id", "text", "similarity")
    .limit(retrieve_k)
    .collect()
)
assert len(retrieved_rows) == retrieve_k
assert [row["document_id"] for row in retrieved_rows] == [
    row["document_id"] for row in expected_rows[:retrieve_k]
]

fallback_ids = [
    row["document_id"]
    for row in sorted(scored_rows, key=lambda row: row["input_position"])[:retrieve_k]
]
assert ANSWER_DOCUMENT_ID not in fallback_ids, fallback_ids
assert retrieved_rows[0]["document_id"] == ANSWER_DOCUMENT_ID, retrieved_rows
assert retrieved_rows[0]["similarity"] > retrieved_rows[1]["similarity"], retrieved_rows

context = "\n\n".join(row["text"] for row in retrieved_rows)
if ci_smoke:
    assert "sparkling blue and white jewel" in retrieved_rows[0]["text"].lower()

spark.createDataFrame(retrieved_rows).show(truncate=False)

6. Generate a grounded answer with Phi-4-mini

Use max_new_tokens, not a fixed total sequence length, so the retrieved context and generated answer cannot conflict. Greedy decoding makes the smoke assertion repeatable. The pinned Phi snapshot is resolved on the single Spark worker and then used as a local path, forcing the model and tokenizer to load the same immutable files. Transformers 4.49 supports this model's native phi3 architecture, so remote model code is disabled. The checkpoint advertises BF16, but a T4 (compute capability 7.5) has no native BF16 support, so model loading explicitly overrides it with FP16.

from synapse.ml.hf import HuggingFaceCausalLM


@pandas_udf(StringType())
def resolve_generation_snapshot(values: pd.Series) -> pd.Series:
    from huggingface_hub import snapshot_download

    snapshot_path = snapshot_download(
        repo_id=GENERATION_MODEL,
        revision=GENERATION_REVISION,
        allow_patterns=["*.json", "*.safetensors", "*.txt"],
    )
    return pd.Series([snapshot_path] * len(values))


generation_snapshot = (
    spark.range(1)
    .repartition(1)
    .select(resolve_generation_snapshot(F.col("id")).alias("snapshot"))
    .first()["snapshot"]
)
assert generation_snapshot.replace("\\", "/").endswith(
    f"/snapshots/{GENERATION_REVISION}"
), f"Unexpected generation snapshot: {generation_snapshot}"

prompt = f"""Use only the context below to answer the question. If the answer is not in the context, say "I don't know."

Context:
{context}

Question: {question}
Answer in one concise sentence."""

prompt_df = spark.createDataFrame([(prompt,)], ["prompt"]).repartition(1)
phi = (
    HuggingFaceCausalLM()
    .setModelName(generation_snapshot)
    .setInputCol("prompt")
    .setOutputCol("answer")
    .setTask("chat")
    .setModelParam(max_new_tokens=generation_tokens, do_sample=False)
    .setModelConfig(
        device_map="cuda",
        torch_dtype="float16",
        local_files_only=True,
        trust_remote_code=False,
    )
)

answer = phi.transform(prompt_df).select("answer").first()["answer"].strip()
assert answer, "Phi returned an empty answer."
if ci_smoke:
    assert "jewel" in answer.lower(), f"Unexpected grounded answer: {answer}"
print(answer)

7. Next steps

This notebook validates the integration seam among Spark GPU UDFs, local vector retrieval, and SynapseML's distributed Hugging Face generation. For production data:

• use the PDF Q&A quickstart's ingestion and chunking stages;
• use an indexed GPU KNN implementation when exact scoring no longer fits the corpus; and
• cache model weights in shared storage as shown in the standalone Phi quickstart.